import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from data_processor import DataProcessor
from model_trainer import ModelTrainer
from database import db  # 导入数据库实例
import os
import joblib
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import (roc_curve, auc, accuracy_score, precision_score, 
                           recall_score, f1_score, confusion_matrix, 
                           classification_report, precision_recall_curve, 
                           average_precision_score, roc_auc_score, make_scorer)
from sklearn.ensemble import RandomForestClassifier
from data_loader import DataLoader
from sklearn.model_selection import cross_validate, StratifiedKFold
from feature_selector import SHAP_RFECV
from xgboost import XGBClassifier


# 设置matplotlib中文显示

os.environ['JOBLIB_TEMP_FOLDER'] = 'D:/joblib_temp'
os.makedirs('D:/joblib_temp', exist_ok=True)
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

# 获取当前文件的目录
BASE_DIR = os.path.dirname(os.path.abspath(__file__))

def evaluate_model(model, X_val, y_val, model_name, is_save_to_file=False):
    """评估单个模型并生成可视化，使用交叉验证"""
    # 使用交叉验证评估模型
    cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
    
    # 定义要计算的评分
    scoring = {
        'accuracy': 'accuracy',
        'precision': 'precision',
        'recall': 'recall',
        'f1': 'f1',
        'roc_auc': 'roc_auc',
        'avg_precision': 'average_precision'
    }
    
    # 进行交叉验证
    cv_results = cross_validate(
        model, X_val, y_val, cv=cv, scoring=scoring, return_train_score=False
    )
    
    # 计算平均指标
    accuracy = np.mean(cv_results['test_accuracy'])
    precision = np.mean(cv_results['test_precision'])
    recall = np.mean(cv_results['test_recall'])
    f1 = np.mean(cv_results['test_f1'])
    roc_auc = np.mean(cv_results['test_roc_auc'])
    avg_precision = np.mean(cv_results['test_avg_precision'])
    
    # 创建评估结果字典
    metrics_dict = {
        'model_name': model_name,
        'accuracy': accuracy,
        'precision': precision,
        'recall': recall,
        'f1': f1,
        'roc_auc': roc_auc,
        'avg_precision': avg_precision
    }
    
    # 如果需要保存到文件
    if is_save_to_file:
        # 创建结果目录
        results_dir = os.path.join(BASE_DIR, 'results')
        os.makedirs(results_dir, exist_ok=True)
        
        # 将结果转换为DataFrame并保存为CSV
        metrics_df = pd.DataFrame([metrics_dict])
        csv_path = os.path.join(results_dir, f'{model_name}_test_evaluation.csv')
        metrics_df.to_csv(csv_path, index=False)
        print(f"\n{model_name}测试集评估结果已保存到: {csv_path}")
        
        # 保存分类报告
        report = classification_report(y_val, model.predict(X_val), output_dict=True)
        report_df = pd.DataFrame(report).transpose()
        report_path = os.path.join(results_dir, f'{model_name}_classification_report.csv')
        report_df.to_csv(report_path)
        print(f"{model_name}分类报告已保存到: {report_path}")
    else:
        # 打印评估结果
        print(f"\n{model_name} 模型评估结果:")
        print(f"准确率: {accuracy:.3f}")
        print(f"精确率: {precision:.3f}")
        print(f"召回率: {recall:.3f}")
        print(f"F1分数: {f1:.3f}")
        print(f"ROC-AUC: {roc_auc:.3f}")
        print(f"平均精确率: {avg_precision:.3f}")
        
        # 生成分类报告
        print("\n分类报告:")
        print(classification_report(y_val, model.predict(X_val)))
    
    # 生成混淆矩阵
    cm = confusion_matrix(y_val, model.predict(X_val))
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
    plt.title(f'{model_name} 混淆矩阵')
    plt.xlabel('预测标签')
    plt.ylabel('真实标签')
    cm_path = os.path.join(BASE_DIR, 'visualizations', f'{model_name}_confusion_matrix.png')
    plt.savefig(cm_path)
    plt.close()
    
    # 生成PR曲线
    precision_curve, recall_curve, _ = precision_recall_curve(y_val, model.predict_proba(X_val)[:, 1])
    plt.figure(figsize=(8, 6))
    plt.plot(recall_curve, precision_curve)
    plt.title(f'{model_name} PR曲线')
    plt.xlabel('召回率')
    plt.ylabel('精确率')
    pr_path = os.path.join(BASE_DIR, 'visualizations', f'{model_name}_pr_curve.png')
    plt.savefig(pr_path)
    plt.close()
    
    return {
        'accuracy': accuracy,
        'precision': precision,
        'recall': recall,
        'f1': f1,
        'roc_auc': roc_auc,
        'avg_precision': avg_precision
    }

def main():
    # 确保visualizations目录存在
    vis_dir = os.path.join(BASE_DIR, 'visualizations')
    os.makedirs(vis_dir, exist_ok=True)
    
    # 创建一个简单的配置对象
    class Config:
        def __init__(self):
            self.data_path = os.path.join(BASE_DIR, 'data', 'heart_disease_full.csv')
            self.target_column = 'target'
    
    config = Config()
    
    # 确保数据库连接
    if db.connection is None:
        db.connect()
    
    # 加载数据
    data_path = os.path.join(BASE_DIR, 'data', 'heart_disease_full.csv')
    df = pd.read_csv(data_path)
    
    # 数据预处理 - 传递 config 参数
    processor = DataProcessor(config)
    X, y = processor.preprocess_data(df)
    
    # 保证X是DataFrame
    if isinstance(X, np.ndarray):
        feature_names = df.drop(columns=[config.target_column]).columns
        X = pd.DataFrame(X, columns=feature_names)
    
    # 只对训练集做特征选择，防止数据泄漏
    X_train, X_temp, y_train, y_temp = train_test_split(X, y, test_size=0.3, random_state=42)
    X_val, X_test, y_val, y_test = train_test_split(X_temp, y_temp, test_size=0.5, random_state=42)

    print("开始SHAP-RFECV特征选择...")
    selector = SHAP_RFECV(XGBClassifier(n_estimators=100, random_state=42), min_features_to_select=5, cv=5)
    selector.fit(X_train, y_train)
    print(f"最优特征数: {len(selector.selected_features_)}")
    print("最优特征子集:", selector.selected_features_)

    # 只保留最优特征
    X_train = selector.transform(X_train)
    X_val = selector.transform(X_val)
    X_test = selector.transform(X_test)
    
    # 使用ModelTrainer进行训练
    trainer = ModelTrainer(config)
    print("\n开始训练模型...")
    print("注意: 模型选择基于加权评分 (30% F1分数 + 30% ROC-AUC + 20% 召回率 + 20% 精确率)，而非仅基于准确率")
    trainer.train(X_train, y_train)
    
    # 获取训练好的模型
    trained_models = trainer.trained_models
    all_metrics = {}
    
    # 评估所有模型
    for name, model in trained_models.items():
        # 英文名到中文名的映射
        name_map = {
            "AdaBoost": "AdaBoost",
            "GradientBoosting": "梯度提升",
            "RandomForest": "随机森林",
            "LogisticRegression": "逻辑回归",
            "SVC": "支持向量机",
            "XGBoost": "XGBoost"
        }
        zh_name = name_map.get(name, name)
        metrics = evaluate_model(model, X_val, y_val, zh_name)
        all_metrics[zh_name] = metrics
    
    # 绘制所有指标的柱状图对比
    metrics_names = ['准确率', '精确率', '召回率', 'F1分数', 'ROC-AUC', '平均精确率']
    metrics_keys = ['accuracy', 'precision', 'recall', 'f1', 'roc_auc', 'avg_precision']
    
    # 准备数据
    model_names = list(all_metrics.keys())
    metrics_data = {name: [all_metrics[name][key] for key in metrics_keys] for name in model_names}
    
    # 设置图形大小和样式
    plt.figure(figsize=(20, 10))
    
    # 设置柱子的位置
    x = np.arange(len(metrics_names))
    width = 0.15
    
    # 绘制柱状图
    for i, name in enumerate(model_names):
        plt.bar(x + i*width, metrics_data[name], width, label=name)
    
    # 设置图形属性
    plt.xlabel('评估指标', fontsize=12)
    plt.ylabel('分数', fontsize=12)
    plt.title('各模型评估指标对比', fontsize=14, pad=20)
    plt.xticks(x + width*2, metrics_names, rotation=45, ha='right')
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.grid(True, axis='y', linestyle='--', alpha=0.7)
    
    # 调整布局，确保所有元素可见
    plt.subplots_adjust(right=0.85, top=0.9, bottom=0.15)
    
    # 保存图形
    metrics_vis_path = os.path.join(vis_dir, 'model_metrics_comparison.png')
    plt.savefig(metrics_vis_path, dpi=300, bbox_inches='tight', facecolor='white')
    plt.close()
    print(f"\n模型评估指标对比图已保存到: {metrics_vis_path}")
    
    # 保存最佳模型
    best_model = trainer.best_model
    model_path = os.path.join(BASE_DIR, 'models', 'best_model.joblib')
    os.makedirs(os.path.dirname(model_path), exist_ok=True)
    joblib.dump(best_model, model_path)
    print(f"\n最佳模型已保存到: {model_path}")
    
    # 在测试集上评估最佳模型
    print("\n在测试集上评估最佳模型:")
    test_metrics = evaluate_model(best_model, X_test, y_test, "最佳模型", is_save_to_file=True)
    
    # 保存测试集评估结果到数据库
    if db.connection is None:
        db.connect()
    db_id = db.save_model_performance("最佳模型", test_metrics)
    print(f"模型评估结果已保存到数据库，ID: {db_id}")
    
    # 绘制ROC曲线
    plt.figure(figsize=(12, 8))
    for name, model in trained_models.items():
        y_pred_proba = model.predict_proba(X_test)[:, 1]
        fpr, tpr, _ = roc_curve(y_test, y_pred_proba)
        roc_auc = auc(fpr, tpr)
        plt.plot(fpr, tpr, label=f'{name} (AUC = {roc_auc:.3f})')
    
    plt.plot([0, 1], [0, 1], 'k--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('假阳性率', fontsize=12)
    plt.ylabel('真阳性率', fontsize=12)
    plt.title('各模型ROC曲线对比', fontsize=14, pad=20)
    plt.legend(loc="lower right", fontsize=10)
    plt.grid(True, linestyle='--', alpha=0.7)
    
    # 保存ROC曲线图
    vis_path = os.path.join(BASE_DIR, 'visualizations', 'roc_curves.png')
    os.makedirs(os.path.dirname(vis_path), exist_ok=True)
    plt.savefig(vis_path, dpi=300, bbox_inches='tight', facecolor='white')
    plt.close()
    print("\nROC曲线已保存到:", vis_path)
    
    # 保存测试集上的混淆矩阵为CSV文件
    y_pred_test = best_model.predict(X_test)
    cm_test = confusion_matrix(y_test, y_pred_test)
    cm_df = pd.DataFrame(cm_test, 
                       index=['实际为阴性', '实际为阳性'], 
                       columns=['预测为阴性', '预测为阳性'])
    cm_csv_path = os.path.join(BASE_DIR, 'results', 'Best model_confusion_matrix.csv')
    cm_df.to_csv(cm_csv_path)
    print(f"测试集混淆矩阵已保存到: {cm_csv_path}")
    
    # 保存详细的测试集指标
    test_metrics_summary = {
        'TP': cm_test[1, 1],  # 真阳性
        'TN': cm_test[0, 0],  # 真阴性
        'FP': cm_test[0, 1],  # 假阳性
        'FN': cm_test[1, 0],  # 假阴性
        'accuracy': accuracy_score(y_test, y_pred_test),
        'precision': precision_score(y_test, y_pred_test),
        'recall': recall_score(y_test, y_pred_test),
        'f1': f1_score(y_test, y_pred_test),
        'avg_precision': average_precision_score(y_test, best_model.predict_proba(X_test)[:, 1])
    }
    test_summary_df = pd.DataFrame([test_metrics_summary])
    test_summary_path = os.path.join(BASE_DIR, 'results', 'Best model_detail.csv')
    test_summary_df.to_csv(test_summary_path, index=False)
    print(f"测试集详细指标已保存到: {test_summary_path}")

if __name__ == "__main__":
    main()